# src/crypto/lfsr.py
class LFSR:
    def __init__(self, seed, taps, num_bits):
        """
        初始化LFSR。
        :param seed: 初始状态（种子），整数形式。
        :param taps: 反馈多项式的位置（列表），例如 [21, 2] 表示 x^21 + x^2 + 1。
        :param num_bits: LFSR的位数。
        """
        if seed == 0:
            raise ValueError("种子不能为全0！")  # 种子为0会导致LFSR无法生成有效序列
        self.seed = seed  # 保存种子
        self.taps = taps  # 保存反馈多项式的位置
        self.num_bits = num_bits  # 保存LFSR的位数
        self.state = seed  # 初始化当前状态为种子

    def shift(self):
        """
        执行一次LFSR移位操作，并返回输出位。
        :return: 输出位（0或1）。
        """
        # 确保 self.state 是整数类型
        if not isinstance(self.state, int):
            raise TypeError(f"self.state 必须是整数类型，当前类型为 {type(self.state)}")

        # 1. 获取输出位：当前状态的最低位
        output_bit = self.state & 1

        # 2. 计算反馈位：
        # 遍历 taps 列表，取出对应位置的比特，求和后取模2
        feedback_bit = sum((self.state >> tap) & 1 for tap in self.taps) % 2

        # 3. 更新状态：
        # 将当前状态右移一位，然后将反馈位左移到最高位，合并为新状态
        self.state = (self.state >> 1) | (feedback_bit << (self.num_bits - 1))

        # 4. 返回输出位
        return output_bit

    def generate_m_sequence(self, length):
        """
        生成指定长度的m序列。
        :param length: 需要生成的序列长度。
        :return: m序列（列表形式）。
        """
        # 调用 shift 方法 length 次，生成指定长度的比特序列
        return [self.shift() for _ in range(length)]

    def generate_key_stream(self, length):
        """
        生成指定长度的密钥流（字节流）。
        :param length: 需要生成的密钥流长度（字节数）。
        :return: 密钥流（字节流）。
        """
        # 1. 生成比特流：调用 generate_m_sequence 方法生成 length * 8 个比特
        m_sequence = self.generate_m_sequence(length * 8)

        # 2. 将比特流转换为字节流：
        # 每8个比特组合成一个字节，使用 int("...", 2) 将二进制字符串转换为整数
        m_sequence_bytes = bytes(int("".join(map(str, m_sequence[i:i + 8])), 2) for i in range(0, len(m_sequence), 8))

        # 3. 返回指定长度的密钥流
        return m_sequence_bytes[:length]
